'''
    The Function of SearchImage(Rubbish English -- something wrong with my input method)
'''

import h5py
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.models as models

import LoadImage as lim

query = 'cat1.jpg' # use this to search
index = 'FeaturesAndPaths/vgg_featureCNN.h5' # location
result = 'DB/pictures'

h5f = h5py.File(index,'r') # open this file.h5
features = h5f['features'][:] # get features
paths = h5f['paths'][:] # get paths
h5f.close() # close this file

# print(features)
# print(paths)

print('------------------------------------------\n'
      '            search image start            \n'
      '------------------------------------------')

# read and show query image
queryImage = mpimg.imread(query)
plt.title("The image I want to find")
plt.imshow(queryImage)
plt.show()

# init VGG16 model (custom dimensionality)
model = models.vgg16(pretrained=False)
pthfile = 'Model/state/vgg16-397923af.pth'  # 预训练模型的地址
model.load_state_dict(torch.load(pthfile))
model.eval()  # 表示用于预测

# Function
# def searchImageFun():
#     img = lim.synthesize(query)
#     queryFeat = model(img)
#     queryFeat = queryFeat.view(queryFeat.size(1))  # transform to one-dimensional
#     # use cosine similarity method
#     scores = np.dot(queryFeat.detach().numpy(), features.T)  # inner product
#     # print(scores)
#     rank_ID = np.argsort(scores)[::-1]  # sort and return index value
#     # print(rank_ID)
#     rank_scores = scores[rank_ID]
#     # print(rank_scores)
#
#     maxres = 3  # number that need to be return
#     imlist = []  # store the most likely pictures
#     for epoch_idx, index in enumerate(rank_ID[0:maxres]):
#         imlist.append(paths[index])
#     return imlist

if __name__ == '__main__':
    img = lim.synthesize(query)
    queryFeat = model(img)
    # queryFeat = SE_VGG(img)
    queryFeat = queryFeat.view(queryFeat.size(1)) # 变换为一维
    # 使用余弦相似法
    scores = np.dot(queryFeat.detach().numpy(),features.T) #向量求点积
    # print(scores)
    rank_ID = np.argsort(scores)[::-1] # 点积值从大到小排列（点积越大向量越相似）
    # print(rank_ID)
    rank_scores = scores[rank_ID]
    # print(rank_scores)

    maxres = 2# 需要返回的数量

    imlist = []#存储最有可能的图片
    for epoch_idx, index in enumerate(rank_ID[0:maxres]):
        imlist.append(paths[index])
        print(f"imagePath : {str(paths[index])} scores = {rank_scores[epoch_idx]}")
        print("top %d images in order are: " %maxres , imlist)

    for idx, path in enumerate(imlist):
        print(path)
        image = mpimg.imread(str(path,'utf-8')) # 获取对应的图片
        plt.title('Searching Results')
        plt.imshow(image)
        plt.show()


